import sys
sys.path.append("/home/fd_chen/in_context_RL/spoc-robot-training")

import os
import json
import shutil
import gc
import time
import random
import math
from tqdm import tqdm
from multiprocessing import Pool
from concurrent.futures import ThreadPoolExecutor

from collections import deque, defaultdict
from matplotlib import pyplot as plt

from environment.stretch_controller import StretchController
from expert.actions import shortest_path_to_actions, plot_grid_map, determine_turn_actions_
from expert.functions import l2_distance, load_procthor_houses, save_data_json, check_agent_position, save_img_from_frame, positions2path
from utils.data_generation_utils.navigation_utils import is_any_object_sufficiently_visible_and_in_center_frame
from utils.constants.stretch_initialization_utils import STRETCH_ENV_ARGS, AGENT_ROTATION_DEG, AGENT_MOVEMENT_CONSTANT

# Constants
MEDIUM_RANGE = (0.2, 0.6)
TARGET_RANGE = 0.8  # All targets
DEBUG = True
MAX_FRAME = 1000
MAX_AGENT_OBJECT_DISTANCE = 1
MAX_RANDOM_AGENT_POSITION_ATTEMPTS = 200
MAX_ATTEMPTS = 200
NUM_PROCESSES = 48
ABSOLUTE_MIN_PIXELS = 100
ALNATIVE_ACTIONS = [ 'm','b', 'l', 'r', 'ls','rs']
# Initialize counters for successful trajectories, their lengths, and object types

def async_save_image(file_path, image):
    """异步保存图像"""
    plt.imsave(file_path, image)

class Expert(StretchController):
    def __init__(self, id=0, **kwargs):
        """Initialize the Expert controller."""
        super().__init__(**kwargs)
        self.id = id

    def check_dir(self,object_id=None):
        """Check and create necessary directories."""
    
        self.frame_dir = f"./generator_data/{self.id}/{object_id}/frame"
        self.metadata_dir = f"./generator_data/{self.id}/{object_id}/metadata"
        
        for dir in [self.frame_dir, self.metadata_dir]:
            if os.path.exists(dir):
                shutil.rmtree(dir)
            os.makedirs(dir, exist_ok=True)
            

    def async_save_navigation_frame(self, nav_frames, frame_dir):
        """异步保存导航帧到目录"""
        self.executor = ThreadPoolExecutor(max_workers=4)  # 设置最大线程数
        for i, nav_frame in enumerate(nav_frames):
            file_path = f"{frame_dir}/{i}.png"
            self.executor.submit(async_save_image, file_path, nav_frame) 
        
        
    def save_navigation_frame(self, nav_frames, frame_dir):
        """Save navigation frames to a directory."""
        for i, nav_frame in enumerate(nav_frames):
            plt.imsave(f"{frame_dir}/{i}.png", nav_frame)
            

    def save_collected_data(self, old_agent_pose, positions, actions, nav_frames):
        """Save agent's old pose, positions, actions, and navigation frames."""
        save_data_json(old_agent_pose, f"{self.metadata_dir}/agent.json")
        save_data_json(positions, f"{self.metadata_dir}/positions.json")
        save_data_json(actions, f"{self.metadata_dir}/actions.json")
        self.async_save_navigation_frame(nav_frames, self.frame_dir)
        
                 
    def filter_objects(self):
        """Filter objects based on target types."""
        return [o for o in self.get_objects_() if o["objectType"] in self.target_types()]

    def target_types(self, datapath="./expert/domain/procthor.yaml"):
        """Load target object types from a YAML file."""
        import yaml
        with open(datapath, 'r') as file:
            data = yaml.safe_load(file)
        return data.get('target_object_types', [])

    def get_objects_(self):
        """Get objects from the environment."""
        with self.include_object_metadata_context():
            return list(self.controller.last_event.metadata["objects"])

    #split the reachable points
    def split_reachable_area(self, object_id):
        """Split the reachable area into difficulty levels."""
        length_from_positions = [(l2_distance(self.get_obj_pos_from_obj_id(object_id), pos), pos) for pos in self.get_reachable_positions()]
        length_from_positions.sort(key=lambda x: x[0])
        return {
            "easy": length_from_positions[
                : int(len(length_from_positions) * MEDIUM_RANGE[0])
            ],
            "medium": length_from_positions[
                int(len(length_from_positions) * MEDIUM_RANGE[0]) : int(
                    len(length_from_positions) * MEDIUM_RANGE[1]
                )
            ],
            "hard": length_from_positions[
                int(len(length_from_positions) * MEDIUM_RANGE[1]) :
            ],
        }

    def random_agent_position(self,objectid,difficulty=random.choice(["easy","medium","hard"])):
        """Set a random agent position."""
        old_agent_location = self.get_current_agent_full_pose()
        self.difficulty_spots =self.split_reachable_area(objectid)[difficulty]
        position = self.difficulty_spots[random.randint(0, len(self.difficulty_spots) - 1)][1]

        rot_y = random.choice(list(range(0, 360, int(AGENT_ROTATION_DEG))))
        
        event = self.controller.step(
            action="Teleport",
            position=position,
            rotation=dict(x=0, y=rot_y, z=0),
            standing=old_agent_location["isStanding"],
        )
        return event.metadata["lastActionSuccess"]

    def teleport_agent(self, agent_pose, teleport_horizon=False):
        """Teleport agent to a given pose."""
        action_params = {
            'action': 'Teleport',
            'position': agent_pose['position'],
            'rotation': agent_pose['rotation'],
            'standing': agent_pose['isStanding']
        }
        if teleport_horizon:
            action_params['horizon'] = agent_pose["cameraHorizon"]
        return self.controller.step(**action_params).metadata["lastActionSuccess"]

    def more_objects_id_name(self):
        """Get PART objects."""
        all_objects_id, all_objects_name,all_objects_info = self.get_all_objects()
        selected_indices = random.sample(range(len(all_objects_id)), int(len(all_objects_id) * TARGET_RANGE))
        selected_objects_id = [all_objects_id[i] for i in selected_indices]
        selected_objects_name = [all_objects_name[i] for i in selected_indices]
        selected_objects_info = [all_objects_info[i] for i in selected_indices]
        return selected_objects_id,selected_objects_info, selected_objects_name


    def get_all_objects(self):
        """Get all object IDs and names."""
        objects = self.filter_objects()
        return [o["objectId"] for o in objects], [o["name"] for o in objects], objects


    def get_actions_from_shortest_path(self, object_id):
        """Convert shortest path to actions."""
        shortest_path = self.get_shortest_path_to_object(object_id)
        return shortest_path_to_actions(
            shortest_path, self.get_current_agent_full_pose()
        )

    def get_agent_distance_to_object(self, object_id):
        """Get the distance between the agent and an object."""
        object_info = self.get_object(object_id)
        agent_info = self.get_current_agent_full_pose()
        return math.hypot(object_info["position"]["x"] - agent_info["position"]["x"],
                          object_info["position"]["z"] - agent_info["position"]["z"])  
    def perform_action_with_retry(self, action, alternative_actions, max_retries=3):
        """
        Perform an action with retries. If the action fails, try alternative actions.

        :param action: The primary action to perform.
        :param alternative_actions: A list of alternative actions to try if the primary action fails.
        :param max_retries: The maximum number of retries for each action.
        :return: True if an action succeeds, False otherwise.
        """
        retries = 0
        while retries < max_retries:
            agent_event = self.agent_step(action)
            if agent_event.metadata["lastActionSuccess"]:
                return True, action
            retries += 1
            # print(f"Action {action} failed. Retrying ({retries}/{max_retries})...")

        # Try alternative actions
        for alt_action in alternative_actions:
            if alt_action == 'b' and random.random() < 0.8:  # Reduce the probability of action 'b'
                continue
            retries = 0
            while retries < max_retries:
                agent_event = self.agent_step(alt_action)
                if agent_event.metadata["lastActionSuccess"]:
                    return True, alt_action
                retries += 1
                # print(f"Alternative action {alt_action} failed. Retrying ({retries}/{max_retries})...")

        return False, None

    
    def is_agent_position_same_shortest_path_start(self, shortest_path,threshold=0.05):
        """Check if the agent's position is close to the start of the shortest path."""
        agent_position = self.get_current_agent_full_pose()["position"]
        start_position = shortest_path[0]
        return l2_distance(agent_position, start_position) < threshold


    def get_position_rotions(self, agent_location):
        """Get positions and rotations from a list of positions."""
        return [
                    agent_location['position']['x'],
                    agent_location['position']['y'],
                    agent_location['position']['z'],
                    agent_location['rotation']['x'],
                    agent_location['rotation']['y'],
                    agent_location['rotation']['z']
                ]

    def verify_get_actions_from_shortest_path(self, object_id):
        """Verify and get actions from the shortest path to an object."""
        
        # from the current agent position to the object  episode start
        old_agent_pose = self.get_current_agent_full_pose()
        positions =[self.get_position_rotions(old_agent_pose)]
        actions = [' ']
        last_actions_success = deque([True], maxlen=MAX_FRAME)  # Last actions success buffer
        nav_frames = deque([self.navigation_camera], maxlen=MAX_FRAME)  # Navigation frame buffer

        attempts = 0 # try to get the shortest path to the object from the current agent position counter
        while attempts < MAX_ATTEMPTS:
            # Get the shortest path and the actions to the target object
            previous_shortest_path = self.get_shortest_path_to_object(object_id,attempt_path_improvement=True)
            if previous_shortest_path is None:
                raise Exception("Failed to retrieve the shortest path. The object is not reachable.") 

            #  using the attempt_path_improvement flag to improve the path ,so to patch the bug of the path
            if not self.is_agent_position_same_shortest_path_start(previous_shortest_path):
                # print("Agent is not at the start of the shortest path.")
                previous_shortest_path.insert(0, self.get_current_agent_full_pose()["position"]) # add the agent position to the start of the shortest path

            previous_actions = shortest_path_to_actions(previous_shortest_path, self.get_current_agent_full_pose())

            # Plot the current path
            plot_grid_map(previous_shortest_path, previous_actions, self.get_current_agent_full_pose(),
                        step_length=AGENT_MOVEMENT_CONSTANT, save_path=f"{self.metadata_dir}/robot_path_{attempts}.jpg")

            # Execute each action in the shortest path
            for action in previous_actions:
                if action == "end":
                    if not is_any_object_sufficiently_visible_and_in_center_frame(
                        self, [object_id], absolute_min_pixels=ABSOLUTE_MIN_PIXELS
                    ):
                        # raise Exception("Action is end , Object is not visible.")
                        break

                    actions.append("end")
                    self.save_collected_data(old_agent_pose, positions, actions, nav_frames)
                    self.teleport_agent(old_agent_pose)
                    return actions, positions,last_actions_success

                # Save the current agent pose before moving
                previous_agent_location = self.get_current_agent_full_pose()
                # Perform the action
                
                # agent_event = self.agent_step(action)
                action_is_success, action = self.perform_action_with_retry(action, alternative_actions=ALNATIVE_ACTIONS)
                if not action_is_success:
                    raise Exception(f"Action {action} and all alternatives failed.")
                
                last_actions_success.append(True) 

                full_last_agent_location = self.get_current_agent_full_pose()

                if self.get_agent_distance_to_object(object_id)<= MAX_AGENT_OBJECT_DISTANCE:
                    if not is_any_object_sufficiently_visible_and_in_center_frame(
                        self, [object_id], absolute_min_pixels=ABSOLUTE_MIN_PIXELS
                    ):
                        # print("Object is not visible.")
                        angle_diff = self.get_agent_alignment_to_object(object_id)
                        actions_ = determine_turn_actions_(angle_diff)
                        for action in actions_:
                            agent_event = self.agent_step(action)
                            if not agent_event.metadata["lastActionSuccess"]:
                                raise Exception(f"Action {action} failed.")
                            last_actions_success.append(True)
                            nav_frames.append(self.navigation_camera)
                            positions.append(self.get_position_rotions(full_last_agent_location))
                            actions.append(action)
                            if is_any_object_sufficiently_visible_and_in_center_frame(
                                self, [object_id], absolute_min_pixels=ABSOLUTE_MIN_PIXELS
                            ):
                                # print("Object is visible.")
                                break
                    actions.append("end")
                    self.save_collected_data(old_agent_pose, positions, actions, nav_frames)
                    self.teleport_agent(old_agent_pose)
                    return actions, positions,last_actions_success

                # Check if the agent deviates from the shortest path
                if not check_agent_position(previous_shortest_path, full_last_agent_location, threshold=0.2):
                    self.teleport_agent(previous_agent_location)
                    last_actions_success.pop()
                    # print("Agent deviates from the shortest path.")
                    break  # Retry from the start since the agent deviated

                # Record frames, positions, and actions
                nav_frames.append(self.navigation_camera)
                positions.append(self.get_position_rotions(full_last_agent_location))
                actions.append(action)

            attempts += 1

        # debugpy.breakpoint()
        return Exception("Max attempts reached. Object is still not visible.")

def process_house(args):

    houses, house_id = args  # Unpack the arguments
    start_time = time.time()  # Start timing the processing of the house
    max_init_attempts = 5
    init_attempts = 0
    STRETCH_ENV_ARGS['gpu_device'] = house_id % 3  # Assign a GPU device to the Expert controller
    while init_attempts < max_init_attempts:
        try:
            Expert_controller = Expert(id=house_id, **STRETCH_ENV_ARGS)
            Expert_controller.reset(houses[house_id])
            break
        except TimeoutError as e:
            init_attempts += 1
            # print(f"Error initializing the Expert controller: {e}. Retrying ({init_attempts}/{max_init_attempts})...")
            time.sleep(2)  # Wait before retrying

    if init_attempts >= max_init_attempts:
        # print(f"Failed to initialize the Expert controller for house {house_id} after {max_init_attempts} attempts.")
        return

    os.makedirs(f"./generator_data/{house_id}", exist_ok=True)
    save_data_json(houses[house_id], f"./generator_data/{house_id}/house.json")
    # print(f"Processing house {house_id}...")
    
    for target_id, target_info in zip(*Expert_controller.more_objects_id_name()[:2]):
        print(target_id)
        RETRIEVE = True
        attemp = 0
        while RETRIEVE and attemp < MAX_RANDOM_AGENT_POSITION_ATTEMPTS:
            try:
                Expert_controller.check_dir(object_id=target_id)
                save_data_json(target_info, f"{Expert_controller.metadata_dir}/object_info.json")

                Expert_controller.random_agent_position(target_id)
                    
                shortest_path = Expert_controller.get_shortest_path_to_object(target_id)
                if shortest_path is None:
                    raise Exception("Failed to retrieve the shortest path.")

                # top_down_path_frame = Expert_controller.get_top_down_path_view(shortest_path)
                # save_img_from_frame(top_down_path_frame, save_path=f"./generator_data/{i}/metadata/test_top_down_shortest_path.jpg")

                actions, positions, _ = Expert_controller.verify_get_actions_from_shortest_path(target_id)

                shortest_path = positions2path(positions)
                top_down_path_frame = Expert_controller.get_top_down_path_view(shortest_path)
                save_img_from_frame(top_down_path_frame, save_path=f"{Expert_controller.metadata_dir}/test_top_down_along_path.jpg")

                init_position = Expert_controller.get_current_agent_full_pose()
                plot_grid_map(shortest_path, actions, init_position, step_length=AGENT_MOVEMENT_CONSTANT, save_path=f"{Expert_controller.metadata_dir}/robot_path.jpg")
                
                print(f"House {house_id} object {target_id} completed.")
                RETRIEVE = False
                gc.collect()
                time.sleep(0.5)

                # Update counters for successful trajectories and object types
                # Record the successful find

            except Exception as e:
                print(f"Error: {e}. Retrying with new agent position...")
                attemp += 1
                RETRIEVE = True  # Keep RETRIEVE
                gc.collect()

    Expert_controller.stop()
    # print(f"House {house_id} completed.")
    gc.collect()

    end_time = time.time()  # End timing the processing of the house
    processing_time = end_time - start_time
    # print(f"Processing house {house_id} took {processing_time:.2f} seconds.")
    
if __name__ == "__main__":
    # enable_remote_debug(65532)
    houses = load_procthor_houses()
    shutil.rmtree("./generator_data", ignore_errors=True)
    # houses = list(load_procthor_houses())[:1]
    with Pool(processes=NUM_PROCESSES) as pool:
        # Use the pool to process each house in parallel with a progress bar
        for _ in tqdm(pool.imap_unordered(process_house, [(houses, house_id) for house_id in range(len(houses))]), total=len(houses)):
            pass

